# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-control-output-arrays=AssignAdd -o - | FileCheck %s
# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-prune-unused-nodes -tf-control-output-arrays=AssignAdd -o - | FileCheck --check-prefix=PRUNE %s
# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-prune-unused-nodes -tf-control-output-arrays=Variable/Assign,AssignAdd -o - | FileCheck --check-prefix=PRESERVE %s

# Generated in Python via
# ```
# import tensorflow as tf
#
# with tf.compat.v1.Graph().as_default() as g:
#   var = tf.Variable(2.0)
#   var_add = var.assign_add(3.0)
# ```

node {
  name: "Variable/initial_value"
  op: "Const"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
        }
        float_val: 2.0
      }
    }
  }
}
node {
  name: "Variable"
  op: "VariableV2"
  attr {
    key: "container"
    value {
      s: ""
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }
  attr {
    key: "shared_name"
    value {
      s: ""
    }
  }
}
node {
  name: "Variable/Assign"
  op: "Assign"
  input: "Variable"
  input: "Variable/initial_value"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@Variable"
      }
    }
  }
  attr {
    key: "use_locking"
    value {
      b: true
    }
  }
  attr {
    key: "validate_shape"
    value {
      b: true
    }
  }
}
node {
  name: "Variable/read"
  op: "Identity"
  input: "Variable"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@Variable"
      }
    }
  }
}
node {
  name: "AssignAdd/value"
  op: "Const"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT
        tensor_shape {
        }
        float_val: 3.0
      }
    }
  }
}
node {
  name: "AssignAdd"
  op: "AssignAdd"
  input: "Variable"
  input: "AssignAdd/value"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@Variable"
      }
    }
  }
  attr {
    key: "use_locking"
    value {
      b: false
    }
  }
}
versions {
  producer: 309
}

# Tests single target node with no pruning set. All nodes will remain in the
# graph and the target node is added to the graph fetch as a control.
#
# CHECK-LABEL: func @main
# CHECK-SAME:  control_outputs = "AssignAdd"
# CHECK-SAME:  inputs = ""
# CHECK-SAME:  outputs = ""
# CHECK:           %{{.*}}, %[[ASSIGN_ADD_CTRL:.*]] = tf_executor.island wraps "tf.AssignAdd"
# CHECK:           tf_executor.fetch %[[ASSIGN_ADD_CTRL]]

# Tests single target node with pruning set. Unreachable nodes from the fetch,
# including side effecting nodes, will be removed. The target node is added to
# the graph fetch as a control.
#
# PRUNE-LABEL: func @main
# PRUNE-SAME:  control_outputs = "AssignAdd"
# PRUNE-SAME:  inputs = ""
# PRUNE-SAME:  outputs = ""
# PRUNE-NOT:       "tf.Assign"
# PRUNE-NOT:       "tf.Identity"
# PRUNE-DAG:       %[[CONST:.*]], %{{.*}} = tf_executor.island wraps "tf.Const"
# PRUNE-DAG:       %[[VAR:.*]], %{{.*}} = tf_executor.island wraps "tf.VariableV2"
# PRUNE:           %{{.*}}, %[[ASSIGN_ADD_CTRL:.*]] = tf_executor.island wraps "tf.AssignAdd"(%[[VAR]], %[[CONST]])
# PRUNE-NEXT:      tf_executor.fetch %[[ASSIGN_ADD_CTRL]]

# Tests multiple target nodes with pruning set. Unreachable nodes from the
# fetch, including side effecting nodes, will be removed. The target nodes are
# added to the graph fetch as controls.
#
# PRESERVE-LABEL: func @main
# PRESERVE-SAME:  control_outputs = "Variable/Assign,AssignAdd"
# PRESERVE-SAME:  inputs = ""
# PRESERVE-SAME:  outputs = ""
# PRESERVE-NOT:       "tf.Identity"
# PRESERVE:           %[[VAR:.*]], %{{.*}} = tf_executor.island wraps "tf.VariableV2"
# PRESERVE-DAG:       %{{.*}}, %[[ASSIGN_CTRL:.*]] = tf_executor.island wraps "tf.Assign"
# PRESERVE-DAG:       %{{.*}}, %[[ASSIGN_ADD_CTRL:.*]] = tf_executor.island wraps "tf.AssignAdd"
# PRESERVE-NEXT:      tf_executor.fetch %[[ASSIGN_CTRL]], %[[ASSIGN_ADD_CTRL]]
